# coding=utf-8
# author : soliton
# gmail  : soliton.wang@gmail.com
# qq     : 1670829014

import re
import sys
import random
import requests
from tabulate import tabulate
from useragent import UserAgent

class Tianyi:
    def __init__(self):
        ua = UserAgent()
        self.session = requests.Session()
        self.Tianyi_url = "http://192.168.1.1/cgi-bin/luci"
        self.pmset = "http://192.168.1.1/cgi-bin/luci/admin/settings/pmSetSingle"
        self.headers = {"User-Agent": ua.random()}

    def login(self):
        """
        :return: 获取token
        """
        data = {
            "username": user,
            "psd": pwd
        }
        login_tianyi = self.session.post(url=self.Tianyi_url, headers=self.headers, data=data).text
        token = re.search(r"token: '(.*)'", login_tianyi).group(1)
        return token

    def add_port(self, name, ip, exPort, inPort, protocol="TCP"):
        """
        :param name: 映射服务名
        :param ip: 端口转换的内网机器IP
        :param exPort: 外网端口
        :param inPort: 内网端口
        :param protocol: 协议：TCP、UDP、BOTH
        :return: 返回成功或失败

        """
        data = {
            "token": self.login(),
            "op": "add",
            "srvname": name,
            "client": ip,
            "protocol": protocol,
            "exPort": exPort,
            "inPort": inPort,
            "_": random.random()
        }

        result = self.session.post(url=self.pmset, headers=self.headers, data=data).json()
        print(result)

        if result['retVal'] == 0:
            print("Successfully added port")
        else:
            print("Failed to add mapping port")

    def all_port_mapping(self, status):
        """
        :param status: enable | disable | del
        :return:
        """
        url = "http://192.168.1.1/cgi-bin/luci/admin/settings/pmSetAll"
        data = {
            "token": self.login(),
            "op": status,
            "_": random.random()
        }
        result = self.session.post(url=url, headers=self.headers, data=data).text
        return result

    def port_mapping(self, name, op):
        """
        :param name: 映射服务名
        :param op: del|disable|enable
        :return: bool
        """
        data = {
            "token": self.login(),
            "op": op,
            "srvname": name,
            "_": random.random()
        }
        result = self.session.post(url=self.pmset, headers=self.headers, data=data).json()
        if result['retVal'] == 0:
            return True
        else:
            return False

    def get_mapping(self):
        global token
        token = self.login()
        url = "http://192.168.1.1/cgi-bin/luci/admin/settings/pmDisplay"

        head = ["Service", "Client IP", "Protocol", "InPort", "ExPort", "状态"]

        params = {
            "_": random.random()
        }
        pm = self.session.get(url=url, headers=self.headers, params=params).json()
        data = []
        num = 1
        while True:
            try:
                t = pm[f'pmRule{num}']
                if t['enable'] == 1:
                    status = "启用"
                else:
                    status = "禁用"
                temp_data = [
                    t['desp'], t['client'], t['protocol'], t['inPort'], t['exPort'], status
                ]
                data.append(temp_data)
                num += 1
            except Exception:
                break
        print(tabulate(data, headers=head, tablefmt="grid"))

    def run(self, params):

        x = {
            "-h": [
                "用法：python3 Gataways.py [选项]... [参数]...",
                "-a, --add   : 添加端口",
                "-all, --all : 映射列表操作, 可以删除、关闭、开启所有,慎用",
                "-h, --help  : 查看帮助",
                "-s, --select: 查询映射端口",
                "-p, --port  : 端口操作",
            ],
            "-a": [
                "-a or --add  SrvName  ClientIP  inPort exPort Protocol"
            ],
            "-all": [
                "-all or --all-port status  (status: enable | disable | del)"
            ],
            "-s": [
                "-s or --select 该选项没有参数"
            ],
            "-p": [
                "-p or --port SrvName status (status: enable | disable | del)"
            ]
        }
        if params[1] == "--help" or params[1] == "-h":
            for text in x['-h']:
                print(text, end='\n')
        elif params[1] == "-a" or params[1] == "--add":
            if len(params) == 7:
                try:
                    self.add_port(name=params[2], ip=params[3], inPort=int(params[4]), exPort=int(params[5]), protocol=params[6])
                except Exception:
                    print("参数错误:")
                    for text in x['-a']:
                        print(text, end='\n')
            else:
                print("参数错误:")
                for text in x['-a']:
                    print(text, end='\n')
        elif params[1] == "-all" or params[1] == "--all-port":
            if params[2] == "enable":
                self.all_port_mapping(status="enable")
            elif params[2] == "disable":
                self.all_port_mapping(status="disable")
            elif params[2] == "del" or params[2] == "delete":
                self.all_port_mapping(status="del")
            else:
                print("参数错误:")
                for text in x['-all']:
                    print(text, end='\n')
        elif params[1] == "-s" or params[1] == "--select":
            self.get_mapping()
        elif params[1] == "-p" or params[1] == "--port":
            if params[3] == "enable":
                self.port_mapping(name=params[2], op="enable")
            elif params[3] == "disable":
                self.port_mapping(name=params[2], op="disable")
            elif params[3] == "del" or params[2] == "delete":
                self.port_mapping(name=params[2], op="del")
            else:
                print("参数错误:")
                for text in x['-p']:
                    print(text, end='\n')
        else:
            for text in x['-h']:
                print(text, end='\n')

if __name__ == '__main__':
    user = "useradmin"
    pwd = "pwd"
    params = sys.argv
    app = Tianyi()
    app.run(params=params)
